{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1fe6e643-9453-4381-9445-bd471685fb96",
   "metadata": {},
   "source": [
    "## Exploring the Twitter Emotion dataset using Autolabel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80110a5b-2b3e-45e2-a2da-f6fa00200dff",
   "metadata": {},
   "source": [
    "#### Setup the API Keys for providers that you want to use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "92993c83-4473-4e05-9510-f543b070c7d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: OPENAI_API_KEY=<key>\n"
     ]
    }
   ],
   "source": [
    "%env OPENAI_API_KEY=<key>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c246f85",
   "metadata": {},
   "source": [
    "#### Install the autolabel library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc181e31",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install 'refuel-autolabel[openai]'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84b014d1-f45c-4479-9acc-0d20870b1786",
   "metadata": {},
   "source": [
    "## Run the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c093fe91-3508-4140-8bd6-217034e3cce6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "from autolabel import LabelingAgent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c93fae0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"config_conll2003.json\") as f:\n",
    "     config = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "acb4a3de-fa84-4b94-b17a-7a6fac892a1d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/f5/gb1rtdtn131956qb38lqdwc40000gn/T/ipykernel_96711/3943655472.py:1: UserWarning: WARNING! logprobs is not default parameter.\n",
      "                logprobs was transferred to model_kwargs.\n",
      "                Please confirm that logprobs is what you intended.\n",
      "  agent = LabelingAgent(config=config)\n",
      "/var/folders/f5/gb1rtdtn131956qb38lqdwc40000gn/T/ipykernel_96711/3943655472.py:1: UserWarning: WARNING! top_logprobs is not default parameter.\n",
      "                top_logprobs was transferred to model_kwargs.\n",
      "                Please confirm that top_logprobs is what you intended.\n",
      "  agent = LabelingAgent(config=config)\n"
     ]
    }
   ],
   "source": [
    "agent = LabelingAgent(config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "92667a39",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5cd8021dcc614e798e53df5b7720a210",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌──────────────────────────┬─────────┐\n",
       "│<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> Total Estimated Cost     </span>│<span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\"> $1.3405 </span>│\n",
       "│<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> Number of Examples       </span>│<span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\"> 2000    </span>│\n",
       "│<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> Average cost per example </span>│<span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\"> $0.0007 </span>│\n",
       "└──────────────────────────┴─────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┌──────────────────────────┬─────────┐\n",
       "│\u001b[1;35m \u001b[0m\u001b[1;35mTotal Estimated Cost    \u001b[0m\u001b[1;35m \u001b[0m│\u001b[1;32m \u001b[0m\u001b[1;32m$1.3405\u001b[0m\u001b[1;32m \u001b[0m│\n",
       "│\u001b[1;35m \u001b[0m\u001b[1;35mNumber of Examples      \u001b[0m\u001b[1;35m \u001b[0m│\u001b[1;32m \u001b[0m\u001b[1;32m2000   \u001b[0m\u001b[1;32m \u001b[0m│\n",
       "│\u001b[1;35m \u001b[0m\u001b[1;35mAverage cost per example\u001b[0m\u001b[1;35m \u001b[0m│\u001b[1;32m \u001b[0m\u001b[1;32m$0.0007\u001b[0m\u001b[1;32m \u001b[0m│\n",
       "└──────────────────────────┴─────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #00ff00; text-decoration-color: #00ff00\">───────────────────────────────────────────────── </span>Prompt Example<span style=\"color: #00ff00; text-decoration-color: #00ff00\"> ──────────────────────────────────────────────────</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[92m───────────────────────────────────────────────── \u001b[0mPrompt Example\u001b[92m ──────────────────────────────────────────────────\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">(</span><span style=\"color: #008000; text-decoration-color: #008000\">'You are an expert at extracting Location, Organization, Person, and Miscellaneous entities from text. Your job is</span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">to extract named entities mentioned in text, grouped by entity type.\\n\\nYou will return the extracted attributes as</span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">a json with the following keys:\\n{\\n    \"Location\": \"A list containing all instances of locations in the input. </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">Output an empty list if no such entities exist.\",\\n    \"Organization\": \"A list containing all instances of all </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">organizations in the input. Output an empty list if no such entities exist.\",\\n    \"Person\": \"A list containing all</span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">instances of all people in the input. Output an empty list if no such entities exist.\",\\n    \"Miscellaneous\": \"A </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">list containing all instances of all miscellaneous entities in the input. Output an empty list if no such entities </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">exist.\"\\n}. Don\\'t output anything other than the final JSON in your response.\\n\\nSome examples with their output </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">answers are provided below:\\n\\nExample: PARIS 1996-12-06\\nOutput:\\n{\"Location\": \"[\\\\\"PARIS\\\\\"]\", \"Organization\": </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">\"[]\", \"Person\": \"[]\", \"Miscellaneous\": \"[]\"}\\n\\nExample: PARIS 1996-12-06\\nOutput:\\n{\"Location\": \"[\\\\\"PARIS\\\\\"]\", </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">\"Organization\": \"[]\", \"Person\": \"[]\", \"Miscellaneous\": \"[]\"}\\n\\nExample: LONDON 1996-12-06\\nOutput:\\n{\"Location\": </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">\"[\\\\\"LONDON\\\\\"]\", \"Organization\": \"[]\", \"Person\": \"[]\", \"Miscellaneous\": \"[]\"}\\n\\nExample: LONDON </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">1996-12-06\\nOutput:\\n{\"Location\": \"[\\\\\"LONDON\\\\\"]\", \"Organization\": \"[]\", \"Person\": \"[]\", \"Miscellaneous\": </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">\"[]\"}\\n\\nExample: LONDON 1996-12-06\\nOutput:\\n{\"Location\": \"[\\\\\"LONDON\\\\\"]\", \"Organization\": \"[]\", \"Person\": \"[]\", </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">\"Miscellaneous\": \"[]\"}\\n\\nNow I want you to label the following example:\\nExample: SOCCER - JAPAN GET LUCKY WIN , </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">CHINA IN SURPRISE DEFEAT .\\nOutput:\\n'</span>, <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'title'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'AnswerFormat'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'description'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'Answer to the provided prompt.'</span>, \n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">'type'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'object'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'properties'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'Location'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'title'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'Location'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'type'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'string'</span><span style=\"font-weight: bold\">}</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'Organization'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'title'</span>: \n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">'Organization'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'type'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'string'</span><span style=\"font-weight: bold\">}</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'Person'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'title'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'Person'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'type'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'string'</span><span style=\"font-weight: bold\">}</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'Miscellaneous'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'title'</span>: \n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">'Miscellaneous'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'type'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'string'</span><span style=\"font-weight: bold\">}}</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'required'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'Location'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'Organization'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'Person'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'Miscellaneous'</span><span style=\"font-weight: bold\">]</span>, \n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">'additionalProperties'</span>: <span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'definitions'</span>: <span style=\"font-weight: bold\">{}})</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m(\u001b[0m\u001b[32m'You are an expert at extracting Location, Organization, Person, and Miscellaneous entities from text. Your job is\u001b[0m\n",
       "\u001b[32mto extract named entities mentioned in text, grouped by entity type.\\n\\nYou will return the extracted attributes as\u001b[0m\n",
       "\u001b[32ma json with the following keys:\\n\u001b[0m\u001b[32m{\u001b[0m\u001b[32m\\n    \"Location\": \"A list containing all instances of locations in the input. \u001b[0m\n",
       "\u001b[32mOutput an empty list if no such entities exist.\",\\n    \"Organization\": \"A list containing all instances of all \u001b[0m\n",
       "\u001b[32morganizations in the input. Output an empty list if no such entities exist.\",\\n    \"Person\": \"A list containing all\u001b[0m\n",
       "\u001b[32minstances of all people in the input. Output an empty list if no such entities exist.\",\\n    \"Miscellaneous\": \"A \u001b[0m\n",
       "\u001b[32mlist containing all instances of all miscellaneous entities in the input. Output an empty list if no such entities \u001b[0m\n",
       "\u001b[32mexist.\"\\n\u001b[0m\u001b[32m}\u001b[0m\u001b[32m. Don\\'t output anything other than the final JSON in your response.\\n\\nSome examples with their output \u001b[0m\n",
       "\u001b[32manswers are provided below:\\n\\nExample: PARIS 1996-12-06\\nOutput:\\n\u001b[0m\u001b[32m{\u001b[0m\u001b[32m\"Location\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m\\\\\"PARIS\\\\\"\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\", \"Organization\": \u001b[0m\n",
       "\u001b[32m\"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\", \"Person\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\", \"Miscellaneous\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\"\u001b[0m\u001b[32m}\u001b[0m\u001b[32m\\n\\nExample: PARIS 1996-12-06\\nOutput:\\n\u001b[0m\u001b[32m{\u001b[0m\u001b[32m\"Location\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m\\\\\"PARIS\\\\\"\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\", \u001b[0m\n",
       "\u001b[32m\"Organization\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\", \"Person\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\", \"Miscellaneous\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\"\u001b[0m\u001b[32m}\u001b[0m\u001b[32m\\n\\nExample: LONDON 1996-12-06\\nOutput:\\n\u001b[0m\u001b[32m{\u001b[0m\u001b[32m\"Location\": \u001b[0m\n",
       "\u001b[32m\"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m\\\\\"LONDON\\\\\"\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\", \"Organization\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\", \"Person\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\", \"Miscellaneous\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\"\u001b[0m\u001b[32m}\u001b[0m\u001b[32m\\n\\nExample: LONDON \u001b[0m\n",
       "\u001b[32m1996-12-06\\nOutput:\\n\u001b[0m\u001b[32m{\u001b[0m\u001b[32m\"Location\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m\\\\\"LONDON\\\\\"\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\", \"Organization\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\", \"Person\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\", \"Miscellaneous\": \u001b[0m\n",
       "\u001b[32m\"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\"\u001b[0m\u001b[32m}\u001b[0m\u001b[32m\\n\\nExample: LONDON 1996-12-06\\nOutput:\\n\u001b[0m\u001b[32m{\u001b[0m\u001b[32m\"Location\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m\\\\\"LONDON\\\\\"\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\", \"Organization\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\", \"Person\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\", \u001b[0m\n",
       "\u001b[32m\"Miscellaneous\": \"\u001b[0m\u001b[32m[\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\"\u001b[0m\u001b[32m}\u001b[0m\u001b[32m\\n\\nNow I want you to label the following example:\\nExample: SOCCER - JAPAN GET LUCKY WIN , \u001b[0m\n",
       "\u001b[32mCHINA IN SURPRISE DEFEAT .\\nOutput:\\n'\u001b[0m, \u001b[1m{\u001b[0m\u001b[32m'title'\u001b[0m: \u001b[32m'AnswerFormat'\u001b[0m, \u001b[32m'description'\u001b[0m: \u001b[32m'Answer to the provided prompt.'\u001b[0m, \n",
       "\u001b[32m'type'\u001b[0m: \u001b[32m'object'\u001b[0m, \u001b[32m'properties'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'Location'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'title'\u001b[0m: \u001b[32m'Location'\u001b[0m, \u001b[32m'type'\u001b[0m: \u001b[32m'string'\u001b[0m\u001b[1m}\u001b[0m, \u001b[32m'Organization'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'title'\u001b[0m: \n",
       "\u001b[32m'Organization'\u001b[0m, \u001b[32m'type'\u001b[0m: \u001b[32m'string'\u001b[0m\u001b[1m}\u001b[0m, \u001b[32m'Person'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'title'\u001b[0m: \u001b[32m'Person'\u001b[0m, \u001b[32m'type'\u001b[0m: \u001b[32m'string'\u001b[0m\u001b[1m}\u001b[0m, \u001b[32m'Miscellaneous'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'title'\u001b[0m: \n",
       "\u001b[32m'Miscellaneous'\u001b[0m, \u001b[32m'type'\u001b[0m: \u001b[32m'string'\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m, \u001b[32m'required'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'Location'\u001b[0m, \u001b[32m'Organization'\u001b[0m, \u001b[32m'Person'\u001b[0m, \u001b[32m'Miscellaneous'\u001b[0m\u001b[1m]\u001b[0m, \n",
       "\u001b[32m'additionalProperties'\u001b[0m: \u001b[3;91mFalse\u001b[0m, \u001b[32m'definitions'\u001b[0m: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #00ff00; text-decoration-color: #00ff00\">───────────────────────────────────────────────────────────────────────────────────────────────────────────────────</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[92m───────────────────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from autolabel import AutolabelDataset\n",
    "\n",
    "ds = AutolabelDataset(\"test.csv\", config=config)\n",
    "agent.plan(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "dd703025-54d8-4349-b0d6-736d2380e966",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b6d9b67d4e8b4a3bac29e59be59173dd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Passed in output_schema: {'title': 'AnswerFormat', 'description': 'Answer to the provided prompt.', 'type': \n",
       "'object', 'properties': {'Location': {'title': 'Location', 'type': 'string'}, 'Organization': {'title': \n",
       "'Organization', 'type': 'string'}, 'Person': {'title': 'Person', 'type': 'string'}, 'Miscellaneous': {'title': \n",
       "'Miscellaneous', 'type': 'string'}}, 'required': ['Location', 'Organization', 'Person', 'Miscellaneous'], \n",
       "'additionalProperties': False, 'definitions': {}}\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Passed in output_schema: {'title': 'AnswerFormat', 'description': 'Answer to the provided prompt.', 'type': \n",
       "'object', 'properties': {'Location': {'title': 'Location', 'type': 'string'}, 'Organization': {'title': \n",
       "'Organization', 'type': 'string'}, 'Person': {'title': 'Person', 'type': 'string'}, 'Miscellaneous': {'title': \n",
       "'Miscellaneous', 'type': 'string'}}, 'required': ['Location', 'Organization', 'Person', 'Miscellaneous'], \n",
       "'additionalProperties': False, 'definitions': {}}\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Actual Cost: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0001</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Actual Cost: \u001b[1;36m0.0001\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                             example   Location Organization  \\\n",
      "0  SOCCER - JAPAN GET LUCKY WIN , CHINA IN SURPRI...  [\"JAPAN\"]           []   \n",
      "\n",
      "      Person Miscellaneous                PersonLocationOrgMiscNER_task_label  \\\n",
      "0  [\"CHINA\"]            []  {'Location': '[\"JAPAN\", \"CHINA\"]', 'Organizati...   \n",
      "\n",
      "       Location_label Organization_label Person_label Miscellaneous_label  \\\n",
      "0  [\"JAPAN\", \"CHINA\"]                 []           []          [\"SOCCER\"]   \n",
      "\n",
      "  PersonLocationOrgMiscNER_task_error  \\\n",
      "0                                None   \n",
      "\n",
      "                PersonLocationOrgMiscNER_task_prompt  \\\n",
      "0  You are an expert at extracting Location, Orga...   \n",
      "\n",
      "   PersonLocationOrgMiscNER_task_successfully_labeled  \\\n",
      "0                                               True    \n",
      "\n",
      "            PersonLocationOrgMiscNER_task_annotation  \n",
      "0  successfully_labeled=True label={'Location': '...  \n"
     ]
    }
   ],
   "source": [
    "ds = agent.run(ds, max_items=1)\n",
    "print(ds.df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d7645ab",
   "metadata": {},
   "source": [
    "### Compute confidence scores\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "5fbc1264",
   "metadata": {},
   "outputs": [],
   "source": [
    "config[\"model\"][\"compute_confidence\"] = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "1998f5e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = LabelingAgent(config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "119e6f22",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "agent.plan(\"test.csv\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63c74705",
   "metadata": {},
   "outputs": [],
   "source": [
    "labels, df, metrics_list = agent.run(\"test.csv\", max_items=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31250a7f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
